#include "xen/renderer/model.hpp"
#include "xen/core.hpp"
#include "xen/renderer/renderer3d.hpp"
#include "xen/application.hpp"

namespace Xen {

ModelLoader::ModelLoader(Texture2DLibrary& library, const std::string& filepath):library_(library) {
    Assimp::Importer importer;
    uint32_t flag = aiProcess_Triangulate|aiProcess_CalcTangentSpace|aiProcess_GenSmoothNormals;
#ifdef ENABLE_ASSIMP_VALIDATION
    flag |= aiProcess_ValidateDataStructure;
#endif
    const aiScene *scene = importer.ReadFile(
            filepath,
            flag
            );
    parent_path_ = filepath.substr(0, filepath.find_last_of('/')) + '/';
    scene_ = scene;

    if(!scene || scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode) {
        XEN_CORE_ASSERT(false, "Assimp Error: " + std::string(importer.GetErrorString()));
    }

    processNode(scene_->mRootNode);
}

void ModelLoader::processNode(const aiNode* node) {
    for (unsigned int i = 0; i < node->mNumMeshes; i++) {
        aiMesh *mesh = scene_->mMeshes[node->mMeshes[i]];
        meshes_.push_back(processMesh(mesh));
    }
    for(unsigned int i = 0; i < node->mNumChildren; i++) {
        processNode(node->mChildren[i]);
    }
}

Ref<Mesh> ModelLoader::processMesh(const aiMesh* mesh) {
    std::vector<Vertex> vertices;
    for (unsigned int i = 0; i < mesh->mNumVertices; i++) {
        Vertex vertex;

        vertex.position.x = mesh->mVertices[i].x;
        vertex.position.y = mesh->mVertices[i].y;
        vertex.position.z = mesh->mVertices[i].z;

        vertex.normal.x = mesh->mNormals[i].x;
        vertex.normal.y = mesh->mNormals[i].y;
        vertex.normal.z = mesh->mNormals[i].z;

        vertex.tangent.x = mesh->mTangents[i].x;
        vertex.tangent.y = mesh->mTangents[i].y;
        vertex.tangent.z = mesh->mTangents[i].z;

        if (mesh->mTextureCoords[0]) {
            vertex.texcoord.x = mesh->mTextureCoords[0][i].x;
            vertex.texcoord.y = mesh->mTextureCoords[0][i].y;
        } else {
            vertex.texcoord = glm::vec2(0.0f, 0.0f);
        }

        vertices.push_back(vertex);
    }

    std::vector<unsigned int> indices;
    for (unsigned int i = 0; i < mesh->mNumFaces; i++) {
        aiFace face = mesh->mFaces[i];
        for (unsigned int j = 0; j < face.mNumIndices; j++) {
            indices.push_back(face.mIndices[j]);
        }
    }

    auto result = Mesh::Create(vertices, indices);

    if (mesh->mMaterialIndex >= 0) {
        aiMaterial *material = scene_->mMaterials[mesh->mMaterialIndex];
        std::vector<Ref<Texture2D>> diffuse_maps  = loadMaterialTextures(material, aiTextureType_DIFFUSE);
        std::vector<Ref<Texture2D>> specular_maps = loadMaterialTextures(material, aiTextureType_SPECULAR);
        std::vector<Ref<Texture2D>> emissive_maps = loadMaterialTextures(material, aiTextureType_EMISSIVE);
        std::vector<Ref<Texture2D>> normal_maps   = loadMaterialTextures(material, aiTextureType_HEIGHT);

        auto mtl = Renderer3D::CreateMaterial();
        if (!diffuse_maps.empty()) {
            mtl->SetTexture("material.diffuse_texture", diffuse_maps[0]);
        } else {
            aiColor3D color;
            material->Get(AI_MATKEY_COLOR_DIFFUSE, color);
            mtl->SetData("material.diffuse", ShaderUniformData::Create(glm::vec3(color.r, color.g, color.b)));
        }
        if (!specular_maps.empty()) {
            mtl->SetTexture("material.specular_texture", specular_maps[0]);
        } else {
            aiColor3D color;
            material->Get(AI_MATKEY_COLOR_SPECULAR, color);
            mtl->SetData("material.specular", ShaderUniformData::Create(glm::vec3(color.r, color.g, color.b)));
        }
        if (!emissive_maps.empty()) {
            mtl->SetTexture("material.emissive_texture", emissive_maps[0]);
        } else {
            aiColor3D color;
            material->Get(AI_MATKEY_COLOR_EMISSIVE, color);
            mtl->SetData("material.emission", ShaderUniformData::Create(glm::vec3(color.r, color.g, color.b)));
        }
        if (!normal_maps.empty()) {
            mtl->SetTexture("material.normal_texture", normal_maps[0]);
        }

        result->SetMaterial(mtl);
    }

    return result;
}

std::vector<Ref<Texture2D>> ModelLoader::loadMaterialTextures(aiMaterial* material, aiTextureType type) {
    std::vector<Ref<Texture2D>> textures;
    for(unsigned int i = 0; i < material->GetTextureCount(type); i++) {
        aiString str;
        material->GetTexture(type, i, &str);
        std::string full_path = parent_path_ + std::string(str.C_Str());
        auto texture = library_.Load(full_path, full_path);
        textures.push_back(texture);
    }
    return textures;
}

}
